import csv
import optparse
import os
import time

import joblib
def list_allfile(path):
    if os.path.exists(path):
        files=os.listdir(path)
    else:
        print('this path not exist')
    for file in files:
        if file.endswith('.pkl'):
            return file


if __name__ == '__main__':
    before = time.time()
    parser = optparse.OptionParser()
    parser.add_option('--in', help='[File path], fasta file, required.', dest='infile', action='store', default='features(2).csv')
    parser.add_option('--in_model', help='[File path], trained model file, required.', dest='model', action='store', default='NB_model.pkl')
    parser.add_option('--out', help='[Directory path], output directory, contains model summary information, required.', dest='output', action='store', default='')
    opts, args = parser.parse_args()
    data = eval(str(opts))
    path = data['infile']
    output = data['output']
    model = data['model']
    # k = open(model+'/kmer.txt')
    # # k = open('model/kmer.txt')
    # for line in k:
    #     K = int(line)
    # GetKmerFile(path, K, output)
    # data = []
    # mkdir(output)
    csv_file = csv.reader(open(path))
    # csv_file = csv.reader(open(output + '/FeatureFile.csv'))
    # mod = list_allfile(model)

    NB = joblib.load(model)
    # NB = joblib.load('NB_model.pkl')
    feature_test = []
    sampleID = []
    count = 0
    for line in csv_file:
        if count != 0:
            feature_test.append(line[0:-1])
            sampleID.append(line[-1])
        count += 1
    group = NB.classes_
    # 预测标签
    y_predict = NB.predict(feature_test)
    print(y_predict)
    # lists = sorted(list(set(y_predict)))
    results = open('results.txt', 'w', encoding='utf-8')
    predict_results_NB = NB.predict_proba(feature_test)
    title = ''
    for list in group:
        title += list + '\t'
    results.write('sample\t'+title+'\n')
    items = ''
    for item in range(len(sampleID)):
        items = ''
        for j in range(len(predict_results_NB[0])):
            items += str(predict_results_NB[item][j])+'\t'
        results.write(sampleID[item]+'\t'+items+'\n')

    print('所需时间：'+str(time.time()-before))


